#include <stdio.h>
#include <stdlib.h>
#include <memory>

#include <string>
#include <vector>

#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <glm/gtc/type_ptr.hpp>

#include "node.h"
#include "objectfactory.h"
#include "model.h"
#include "glslprogram.h"
#include "texture.h"

enum VIEW_ANGLE {
    FIRST_VIEW,
    SECOND_VIEW,
};

static float skyboxVertices[] = {
    // positions
    -1.0f*10.0f,  1.0f*10.0f, -1.0f*10.0f, -1.0f,  1.0f, -1.0f,
    -1.0f*10.0f, -1.0f*10.0f, -1.0f*10.0f, -1.0f, -1.0f, -1.0f,
     1.0f*10.0f, -1.0f*10.0f, -1.0f*10.0f,  1.0f, -1.0f, -1.0f,
     1.0f*10.0f, -1.0f*10.0f, -1.0f*10.0f,  1.0f, -1.0f, -1.0f,
     1.0f*10.0f,  1.0f*10.0f, -1.0f*10.0f,  1.0f,  1.0f, -1.0f,
    -1.0f*10.0f,  1.0f*10.0f, -1.0f*10.0f, -1.0f,  1.0f, -1.0f,

    -1.0f*10.0f, -1.0f*10.0f,  1.0f*10.0f, -1.0f, -1.0f,  1.0f,
    -1.0f*10.0f, -1.0f*10.0f, -1.0f*10.0f, -1.0f, -1.0f, -1.0f,
    -1.0f*10.0f,  1.0f*10.0f, -1.0f*10.0f, -1.0f,  1.0f, -1.0f,
    -1.0f*10.0f,  1.0f*10.0f, -1.0f*10.0f, -1.0f,  1.0f, -1.0f,
    -1.0f*10.0f,  1.0f*10.0f,  1.0f*10.0f, -1.0f,  1.0f,  1.0f,
    -1.0f*10.0f, -1.0f*10.0f,  1.0f*10.0f, -1.0f, -1.0f,  1.0f,

     1.0f*10.0f, -1.0f*10.0f, -1.0f*10.0f,  1.0f, -1.0f, -1.0f,
     1.0f*10.0f, -1.0f*10.0f,  1.0f*10.0f,  1.0f, -1.0f,  1.0f,
     1.0f*10.0f,  1.0f*10.0f,  1.0f*10.0f,  1.0f,  1.0f,  1.0f,
     1.0f*10.0f,  1.0f*10.0f,  1.0f*10.0f,  1.0f,  1.0f,  1.0f,
     1.0f*10.0f,  1.0f*10.0f, -1.0f*10.0f,  1.0f,  1.0f, -1.0f,
     1.0f*10.0f, -1.0f*10.0f, -1.0f*10.0f,  1.0f, -1.0f, -1.0f,

    -1.0f*10.0f, -1.0f*10.0f,  1.0f*10.0f, -1.0f, -1.0f,  1.0f,
    -1.0f*10.0f,  1.0f*10.0f,  1.0f*10.0f, -1.0f,  1.0f,  1.0f,
     1.0f*10.0f,  1.0f*10.0f,  1.0f*10.0f,  1.0f,  1.0f,  1.0f,
     1.0f*10.0f,  1.0f*10.0f,  1.0f*10.0f,  1.0f,  1.0f,  1.0f,
     1.0f*10.0f, -1.0f*10.0f,  1.0f*10.0f,  1.0f, -1.0f,  1.0f,
    -1.0f*10.0f, -1.0f*10.0f,  1.0f*10.0f, -1.0f, -1.0f,  1.0f,

    -1.0f*10.0f,  1.0f*10.0f, -1.0f*10.0f, -1.0f,  1.0f, -1.0f,
     1.0f*10.0f,  1.0f*10.0f, -1.0f*10.0f,  1.0f,  1.0f, -1.0f,
     1.0f*10.0f,  1.0f*10.0f,  1.0f*10.0f,  1.0f,  1.0f,  1.0f,
     1.0f*10.0f,  1.0f*10.0f,  1.0f*10.0f,  1.0f,  1.0f,  1.0f,
    -1.0f*10.0f,  1.0f*10.0f,  1.0f*10.0f, -1.0f,  1.0f,  1.0f,
    -1.0f*10.0f,  1.0f*10.0f, -1.0f*10.0f, -1.0f,  1.0f, -1.0f,

    -1.0f*10.0f, -1.0f*10.0f, -1.0f*10.0f, -1.0f, -1.0f, -1.0f,
    -1.0f*10.0f, -1.0f*10.0f,  1.0f*10.0f, -1.0f, -1.0f,  1.0f,
     1.0f*10.0f, -1.0f*10.0f, -1.0f*10.0f,  1.0f, -1.0f, -1.0f,
     1.0f*10.0f, -1.0f*10.0f, -1.0f*10.0f,  1.0f, -1.0f, -1.0f,
    -1.0f*10.0f, -1.0f*10.0f,  1.0f*10.0f, -1.0f, -1.0f,  1.0f,
     1.0f*10.0f, -1.0f*10.0f,  1.0f*10.0f,  1.0f, -1.0f,  1.0f
};

class Computeshader : public Node
{
public:
    Computeshader()
    {
        nodeName = "Computeshader";
    }
public:
    const GLuint WORK_GROUP_X_NUM = 360;
    const GLuint WORK_GROUP_Y_NUM = 100;
    const GLuint WORK_ITEM_X_NUM = 300;
    const GLuint LESS_WORK_GROUP_X_NUM = 36;
    const GLuint LESS_WORK_ITEM_X_NUM = 1;
protected:
    bool startup() override;
    void render(double currentTime, double difTime) override;
    void shutdown() override;
    void onResized(int width, int height)override;

private:
    bool checkSupportGroupSize();
    bool makeSSBO();
    void makeSky();
    void makeFireVao();
    void makeAster();
    void renderFire();
private:
    GLuint ssbo = 0;
    GLSLProgram computeProgram;
    GLSLProgram computeLessProgram;

    GLuint skyVao = 0;
    GLuint skyVbo = 0;
    GLSLProgram skyProgram;
    GLuint skyTex = 0;

    GLuint fireVao = 0;
    GLSLProgram fireProgram;

    std::unique_ptr<Model> asterModel;
    GLSLProgram asterProgram;

    glm::vec3 pos = glm::vec3(0.0f, 0.0f, 0.0f);
    float fireSpeed = 0.000001f;
    float fireStep = 0.000002f;
    const float fireMaxSpeed = 0.0005f;
    glm::mat4 asterTrans = glm::mat4(1.0f);
    glm::mat4 fireTrans = glm::mat4(1.0f);
    glm::mat4 perspec = glm::mat4(1.0f);
    glm::mat4 view = glm::mat4(1.0f);
    glm::vec3 viewPos = glm::vec3(0.014973f, -0.040693f, 0.0924571);
    glm::vec3 viewZAis = glm::vec3(-0.111635f, 0.286525f, -0.843972f);
    glm::vec3 viewYAxis = glm::vec3(0, 1, 0);
    VIEW_ANGLE viewAngle = VIEW_ANGLE::SECOND_VIEW;
};

bool Computeshader::startup()
{
    if (PerfWindow::glVersion < 430) {
        printf("Failed! The GL version is less than 4.3 !\n");
        return false;
    }
    pos = glm::vec3(0.0f, 0.0f, 0.0f);
    fireSpeed = 0.000001f;
    fireStep = 0.000002f;
    asterTrans = glm::mat4(1.0f);
    fireTrans = glm::mat4(1.0f);
    perspec = glm::mat4(1.0f);
    view = glm::mat4(1.0f);
    viewPos = glm::vec3(0.014973f, -0.040693f, 0.0924571);
    viewZAis = glm::vec3(-0.111635f, 0.286525f, -0.843972f);
    viewYAxis = glm::vec3(0, 1, 0);
    if (viewAngle == VIEW_ANGLE::FIRST_VIEW) {
        viewPos = glm::vec3(-0.0538087f, -0.464297f, 0.181731f);
        viewZAis = glm::vec3(0.127917f, 0.207424f, -0.536437f);
    } else if (viewAngle == VIEW_ANGLE::SECOND_VIEW) {
        viewPos = glm::vec3(-0.0527043f, -0.275424f, 0.00421948f);
        viewZAis = glm::vec3(0.587324f, 0.492862f, 0.0142725f);
    }

    perspec = glm::perspective(glm::radians(45.0f), \
                               static_cast<float>(getWidth()) / static_cast<float>(getHeight()), 0.01f, 10000.0f);
    view = glm::lookAt(viewPos, viewZAis, viewYAxis);

    if (!computeProgram.build("../media/shaders/computeshader/cs.glsl") \
            || !computeLessProgram.build("../media/shaders/computeshader/csless.glsl")) {
        return false;
    }
    if (!checkSupportGroupSize()) {
        return false;
    }

    if ( !skyProgram.build("../media/shaders/computeshader/sky.vs.glsl", \
                           "../media/shaders/computeshader/sky.fs.glsl") ) {
        return false;
    }
    if (!fireProgram.build("../media/shaders/computeshader/vs.glsl",
                           "../media/shaders/computeshader/fs.glsl")) {
        return false;
    }

    if (!makeSSBO()) {
        computeProgram.clear();
        skyProgram.clear();
        fireProgram.clear();
        return false;
    }
    makeSky();
    makeFireVao();
    makeAster();

    return true;
}

void Computeshader::render(double currentTime, double difTime)
{
    (void)currentTime;
    glClearColor(0.2f, 0.3f, 0.3f, 0.0f);
    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
    glEnable(GL_DEPTH_TEST);
    glEnable(GL_BLEND);
    glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

    computeProgram.use();
    computeProgram.setUniform1f("speedAdded", fireSpeed);
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, ssbo);
    glDispatchCompute(WORK_GROUP_X_NUM, WORK_GROUP_Y_NUM, 1);
    glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, 0);

    skyProgram.use();
    skyProgram.setUniformMatrix4fv("view", glm::value_ptr(view) );
    glBindTexture(GL_TEXTURE_CUBE_MAP, skyTex);
    glBindVertexArray(skyVao);
    glDrawArrays(GL_TRIANGLES, 0, 36);

    GLfloat movePosY = (-1.0f * fireSpeed * 200.0f * static_cast<GLfloat>(difTime));
    pos = pos + glm::vec3(0.0f, movePosY, 0.0f);

    asterProgram.use();
    asterTrans = glm::mat4(1.0f);
    asterTrans = glm::translate(asterTrans, pos);
    asterTrans = glm::scale(asterTrans, glm::vec3(0.0241f, 0.0241f, 0.0241f));
    asterProgram.setUniformMatrix4fv("trans", glm::value_ptr(asterTrans));
    asterProgram.setUniformMatrix4fv("view", glm::value_ptr(view));
    asterModel->renderModel();

    renderFire();

    fireSpeed = fireSpeed + fireStep * static_cast<GLfloat>(difTime) * 100.0f;
    if (fireSpeed >= fireMaxSpeed) {
        fireSpeed = fireMaxSpeed;
    }

    if (pos.y < -0.4f) {
        pos = glm::vec3(0.0f, 0.0f, 0.0f);
        fireSpeed = 0.000001f;
    }
}

void Computeshader::shutdown()
{
    glDisable(GL_DEPTH_TEST);
    glDisable(GL_BLEND);

    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, 0);
    glBindVertexArray(0);
    glBindTexture(GL_TEXTURE_CUBE_MAP, 0);

    glDeleteBuffers(1, &ssbo);
    glDeleteBuffers(1, &skyVbo);
    glDeleteVertexArrays(1, &skyVao);
    glDeleteVertexArrays(1, &fireVao);
    glDeleteTextures(1, &skyTex);
    computeProgram.clear();
    skyProgram.clear();
    fireProgram.clear();
    asterModel->clearModel();
    asterProgram.clear();
    computeLessProgram.clear();
}

void Computeshader::onResized(int width, int height)
{
    glViewport(0, 0, width, height);
    Node::onResized(width, height);
}

bool Computeshader::checkSupportGroupSize()
{
    computeProgram.use();
    glm::ivec3 gMaxCount = glm::ivec3(0);
    int gInvoMaxSize = 0;
    glGetIntegeri_v(GL_MAX_COMPUTE_WORK_GROUP_COUNT, 0, &gMaxCount[0]);
    glGetIntegeri_v(GL_MAX_COMPUTE_WORK_GROUP_COUNT, 1, &gMaxCount[1]);
    glGetIntegeri_v(GL_MAX_COMPUTE_WORK_GROUP_COUNT, 2, &gMaxCount[2]);
    glGetIntegerv(GL_MAX_COMPUTE_WORK_GROUP_INVOCATIONS, &gInvoMaxSize);
    if ((gMaxCount.x < static_cast<GLint>(WORK_GROUP_X_NUM)  \
            || gMaxCount.y < static_cast<GLint>(WORK_GROUP_Y_NUM)) \
            || gInvoMaxSize < static_cast<GLint>(WORK_ITEM_X_NUM)) {
        return false;
    }
    return true;
}

bool Computeshader::makeSSBO()
{
    glGenBuffers(1, &ssbo);
    glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssbo);
    GLuint errorNum = glGetError();
    GLsizei ssboSize = static_cast<GLint>( (sizeof(glm::vec3) + sizeof(glm::vec4)) \
                                           * WORK_GROUP_X_NUM * WORK_GROUP_Y_NUM * WORK_ITEM_X_NUM ) ;
    glBufferData(GL_SHADER_STORAGE_BUFFER, ssboSize, nullptr, GL_DYNAMIC_DRAW);

    errorNum = glGetError();
    if (errorNum != GL_NO_ERROR) {
        if (errorNum == GL_OUT_OF_MEMORY) {
            std::cout << "gl error : " << " ssbo out of memory" << std::endl;
        }
        std::cout << "gl error : " << errorNum << std::endl;
        glDeleteBuffers(1, &ssbo);
        return false;
    }

    glBindBuffer(GL_SHADER_STORAGE_BUFFER, 0);
    return true;
}

void Computeshader::makeSky()
{
    glGenVertexArrays(1, &skyVao);
    glBindVertexArray(skyVao);
    glGenBuffers(1, &skyVbo);
    glBindBuffer(GL_ARRAY_BUFFER, skyVbo);
    glBufferData(GL_ARRAY_BUFFER, sizeof(skyboxVertices), skyboxVertices, GL_STATIC_DRAW);
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 6 * sizeof(GLfloat), nullptr);
    glEnableVertexAttribArray(0);
    glBindBuffer(GL_ARRAY_BUFFER, 0);
    glBindVertexArray(0);
    skyProgram.use();
    skyProgram.setUniformMatrix4fv("view", glm::value_ptr(view) );
    skyProgram.setUniformMatrix4fv("projection", glm::value_ptr(perspec) );
    std::array<std::string, 6> texFilenames;
    texFilenames[0] = "../media/textures/skybox/right.jpg";
    texFilenames[1] = "../media/textures/skybox/left.jpg";
    texFilenames[2] = "../media/textures/skybox/top.jpg";
    texFilenames[3] = "../media/textures/skybox/bottom.jpg";
    texFilenames[4] = "../media/textures/skybox/front.jpg";
    texFilenames[5] = "../media/textures/skybox/back.jpg";
    Texture::loadCubeMap(texFilenames, &skyTex);
}

void Computeshader::makeFireVao()
{
    glGenVertexArrays(1, &fireVao);
    glBindVertexArray(fireVao);
    glBindBuffer(GL_ARRAY_BUFFER, ssbo);
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 7 * sizeof(GLfloat), nullptr);
    glVertexAttribPointer(1, 4, GL_FLOAT, GL_FALSE, 7 * sizeof(GLfloat), \
                          reinterpret_cast<void *>(3 * sizeof(GLfloat)) );
    glEnableVertexAttribArray(0);
    glEnableVertexAttribArray(1);
    glBindVertexArray(0);
    glBindBuffer(GL_ARRAY_BUFFER, 0);
    fireProgram.use();
    fireTrans = glm::translate(fireTrans, glm::vec3(0.0f, -0.016f, 0.0f));
    fireProgram.setUniformMatrix4fv("model", glm::value_ptr(fireTrans));
    fireProgram.setUniformMatrix4fv("view", glm::value_ptr(view));
    fireProgram.setUniformMatrix4fv("projection", glm::value_ptr(perspec));
}

void Computeshader::makeAster()
{
    asterModel.reset(new Model);
    asterModel->loadModel("../media/objects/planet.obj");
    asterModel->setUseGLPath(GL_FALSE);
    asterProgram.build("../media/shaders/computeshader/rock.vs.glsl", \
                       "../media/shaders/computeshader/rock.fs.glsl");
    asterProgram.use();
    asterProgram.setUniformMatrix4fv("trans", glm::value_ptr(asterTrans));
    asterProgram.setUniformMatrix4fv("view", glm::value_ptr(view));
    asterProgram.setUniformMatrix4fv("projection", glm::value_ptr(perspec));
}

void Computeshader::renderFire()
{
    computeLessProgram.use();
    computeLessProgram.setUniform1f("speedAdded", fireSpeed);
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, ssbo);
    glDispatchCompute(LESS_WORK_GROUP_X_NUM, WORK_GROUP_Y_NUM, 1);
    glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
    glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, 0);

    fireProgram.use();
    fireTrans = glm::mat4(1.0f);
    fireTrans = glm::translate(fireTrans, glm::vec3(pos.x, pos.y - 0.0251f, pos.z));
    fireProgram.setUniformMatrix4fv("model", glm::value_ptr(fireTrans));
    fireProgram.setUniformMatrix4fv("view", glm::value_ptr(view));
    glBindVertexArray(fireVao);
    glDrawArrays(GL_TRIANGLE_FAN, 0, static_cast<GLint>(LESS_WORK_GROUP_X_NUM \
                                                        * WORK_GROUP_Y_NUM \
                                                        * LESS_WORK_ITEM_X_NUM));
}

REGISTER_OBJECT(Node, Computeshader)
